#include "utils.h"
#include <sys/stat.h>
#include <sys/types.h>
#include <fcntl.h>
#include <immintrin.h>
using namespace std;

matrix A, B;
int N;

void multiply_naive(const int n, matrix a, matrix b, matrix c) {
    rep(i,0,n){
        rep(j,0,n){
            double sum = 0;
            rep(k,0,n)
                    sum += a[i][k] * b[k][j];
            c[i][j]=sum;
        }
    }
}

void multiply_cpu(const int N, const int blockSize, const char* fileA, const char* fileB, const char* fileC) {
    matrix a, b, c;
    allocate_matrix(blockSize, a);
    allocate_matrix(blockSize, b);
    allocate_matrix(blockSize, c);
    size_t line_size = blockSize*sizeof(double);
    int fdA = open(fileA, O_RDONLY);
    int fdB = open(fileB, O_RDONLY);
    int fdC = open(fileC, O_CREAT|O_RDWR, 0664);
    timeval start, end;
    gettimeofday(&start, NULL);
    for(int ii=0;ii<N;ii+=blockSize)
        for(int jj=0;jj<N;jj+=blockSize){
            memset_zero(blockSize, c);
            for(int kk=0;kk<N;kk+=blockSize){
                for(int i=0;i<blockSize;++i)
                    pread(fdA, a[i], line_size, ((ii+i)*N+kk)*sizeof(double));
                for(int k=0;k<blockSize;++k)
                    pread(fdB, b[k], line_size, ((kk+k)*N+jj)*sizeof(double));
                for(int i=0;i<blockSize;++i)
                    #pragma omp parallel for
                    for(int j=0;j<blockSize;j+=256/64){

                        __m256d cc = _mm256_load_pd(c[i]+j);
                        for(int k=0;k<blockSize;++k){

                            __m256d aik = _mm256_set1_pd(a[i][k]);
                            __m256d bb = _mm256_load_pd(b[k]+j);
                            cc = _mm256_fmadd_pd(aik, bb, cc);
                        }
                        _mm256_store_pd(c[i]+j, cc);
                    }
                for(int i=0;i<blockSize;++i)
                    pwrite(fdC, c[i], line_size, ((ii+i)*N+jj)*sizeof(double));
            }
        }
    gettimeofday(&end, NULL);
    printf("%.5lfsec consumed\n",end.tv_sec-start.tv_sec+(end.tv_usec-start.tv_usec)*1e-6);

    close(fdA);
    close(fdB);
    close(fdC);
    destroy_matrix(blockSize, a);
    destroy_matrix(blockSize, b);
    destroy_matrix(blockSize, c);
}


void multiply_gpu(const int N, double* a, double* b, double* c);

void init_work() {
    allocate_matrix(N, A);
    allocate_matrix(N, B);
    generate_matrix(N, A);
    generate_matrix(N, B);
}
void end_work() {
    destroy_matrix(N, A);
    destroy_matrix(N, B);
}

void usage(const char* progname) {
    printf("Usage:%s <maxtrix_size> [-c|-g|-b]\n", progname);
    printf("Program Options:\n");
    printf("  -i run with naive method, output to result_naive.dat\n");
    printf("  -g run with cuda method, output to result_gpu.dat\n");
    printf("  -c run with openmp and SIMD, using out memory, output to result_cpu.dat\n");
    printf("  -h for help, this message\n");
    printf("  no option for generating matrix A and B, output to matrixA.dat and matrixB.dat\n");
}

int main(int argc, const char* argv[]){
    N = atoi(argv[1]);
    bool flag_naive = 0;
    bool flag_big = 0;
    bool flag_test = 0;
    bool flag_gpu = 0;
    char *inputA = "matrixA.dat";
    char *inputB = "matrixB.dat";
    char *output_naive = "result_naive.dat";
    char *output_hpc = "result_hpc.dat";
    char *output_gpu = "result_gpu.dat";

    rep(i,1,argc)
    if(strcmp("-i",argv[i])==0) {
        flag_naive = 1;
    } else if(strcmp("-c",argv[i])==0) {
        flag_big = 1;
    } else if(strcmp("-g",argv[i])==0) {
        flag_gpu = 1;
    } else if(strcmp("-t",argv[i])==0) {
        flag_test = 1;
    } else if(strcmp("-h",argv[i])==0) {
        usage(argv[0]);
        return 0;
    }
    printf("%d\n",N);
     if (flag_test) {
        init_work();
        input(N, A, output_naive);
        input(N, B, output_gpu);
        for(int i=0;i<N;++i)
            for(int j=0;j<N;++j)
                if (fabs(A[i][j]-B[i][j])>1e-7)
                    printf("%d %d %.10lf",i,j,A[i][j]-B[i][j]);

        end_work();
    } else if (flag_gpu){
        double *a, *b, *c;
        allocate_matrix(N, a);
        allocate_matrix(N, b);
        allocate_matrix(N, c);
        c = (double*)malloc(sizeof(double)*N*N);
        input(N,a,inputA);
        input(N,b,inputB);
        multiply_gpu(N, a, b, c);
//        for(int i=0;i<N;i++){
//            for(int j=0;j<N;j++)
//                printf("%.5lf\t",c[i*N+j]);
//            printf("\n");
//        }
        output(N,c,output_gpu);
        destroy_matrix(a);
        destroy_matrix(b);
        destroy_matrix(c);
    }
    else if(flag_big){
        multiply_cpu(N, 256, inputA, inputB, output_hpc);
    } else if(flag_naive) {
        init_work();
        matrix C;
        allocate_matrix(N, C);
        input(N, A, inputA);
        input(N, B, inputB);
        test_time(multiply_naive(N,A,B,C));
        output(N, C, output_naive);
//        for(int i=0;i<N;++i){
//            for(int j=0;j<N;++j)
//                printf("%.5lf\t",C[i][j]);
//            printf("\n");
//        }
        destroy_matrix(N, C);
        end_work();
    } else {
         init_work();
         output(N, A, inputA);
         output(N, B, inputB);
         end_work();
     }
    return 0;
}
